import os
import db
import json
import math
import traceback
from datetime import datetime
from utils.format import d2id
from utils.handle import kill_pid
from tqdm import tqdm
from controllers.stock import check_one_strategy_trigger
from settings import DEFAULT_CAL_EXCHANGE
from collections import defaultdict
from multiprocessing import Process
from constants import BACKTESTSTATUS

def income_times_run(
    args,
    strategy_name,
    from_date_int,
    low_price,
    high_price,
    days,
    percent
):
    """策略的收益次数统计
    params:
        args: 总参数
        stratey_name: 策略名称
        from_date_int: 回测的开始时间
        low_price: 筛选最底的股票值
        high_price: 筛选最高的股票值
        days: 统计策略触发后多少天内的股票
        percent: 目标收益率
    """
    # 查找股票
    stocks = db.StockTB.find(name=("not like", "%ST%"), market="主板")
    from_trade_date = db.TradeCalTB.next_trade_day(from_date_int)
    to_trade_date = db.TradeCalTB.next_trade_day()
    strategy = db.StrategyTB.find_one(name=strategy_name)
    if not strategy:
        raise Exception("不存在此策略")
    if strategy["orient"] != 1:
        raise Exception("此回测必须为买入策略")
    today = d2id(datetime.now())

    # 提前取出日历，加速计算
    trade_cals = db.TradeCalTB.find(
            pretrade_date=("between", [from_trade_date, to_trade_date]),
            is_open=1,
            exchange=DEFAULT_CAL_EXCHANGE)
    next_trade_date = {i["pretrade_date"]:i["cal_date"] for i in trade_cals}

    # 创建一个新的回测记录
    backtest_id = db.Backtest.add(
            strategy_name=strategy_name,
            params=json.dumps(args),
            pid=os.getpid(),
            status=BACKTESTSTATUS.RUNNING,
            cal_date=today)

    try:
        # 统计数据：是否上涨到目标，上涨的最大幅值，下降的最大幅值
        stocks_income_times = defaultdict(list)
        for num, stock in enumerate(stocks):
            print("进度：%s/%s" % (num, len(stocks)))
            backtest_info = db.Backtest.find_one(id=backtest_id)
            if backtest_info["status"] != BACKTESTSTATUS.RUNNING:
                break
            db.Backtest.set(percent=math.ceil((num+1)*100/len(stocks))).where(id=backtest_id)
            daily_price = db.DailyTB.find_one(trade_date=from_trade_date, ts_code=stock["ts_code"])
            if (daily_price and
                daily_price.get("close") and
                daily_price["close"] > low_price and
                daily_price["close"] < high_price
            ):
                cal_date_int = from_trade_date
                while cal_date_int <= today:
                    # 判断是否触发
                    try:
                        if check_one_strategy_trigger(stock["ts_code"], strategy, cal_date_int):
                            daily_price = db.DailyTB.find_one(
                                    trade_date=cal_date_int,
                                    ts_code=stock["ts_code"])
                            compare_date_int = cal_date_int

                            # 统计触发后days天后的收益情况
                            max_per, min_per, is_right = 0, 0, 0
                            for i in range(days):
                                if compare_date_int > today:
                                    break
                                compare_date_int = next_trade_date[compare_date_int]
                                compare_daily_price = db.DailyTB.find_one(
                                        trade_date=compare_date_int,
                                        ts_code=stock["ts_code"])
                                if compare_daily_price and compare_daily_price.get("close"):
                                    compare_percent = round(
                                            ((compare_daily_price["close"]-daily_price["close"])*100/daily_price["close"]),
                                            3)
                                    if compare_percent > max_per:
                                        max_per = compare_percent
                                    if compare_percent < min_per:
                                        min_per = compare_percent

                            if max_per > percent:
                                is_right = 1
                            stocks_income_times[str([stock["ts_code"], stock["name"]])].append(dict(
                                day = cal_date_int,
                                price = daily_price["close"],
                                right = is_right,
                                max = max_per,
                                min = min_per
                            ))
                    except Exception as e:
                        print(e)
                    if cal_date_int > today:
                        break
                    else:
                        cal_date_int = next_trade_date[cal_date_int]

        all_times, all_success_times = 0, 0
        for stock_info, income_lst in stocks_income_times.items():
            times = len(income_lst)
            success_times = len([i for i in income_lst if i["right"]])
            db.BacktestRecord.add(
                backtest_id=backtest_id,
                stock_info=stock_info,
                key1='触发次数',
                value1=str(times),
                key2='达到目标次数',
                value2=str(success_times),
                mark=str(income_lst),
                cal_date=today
            )
            all_times += times
            all_success_times += success_times

        res = dict(
            key1='触发次数',
            value1=str(all_times),
            key2='达到目标次数',
            value2=str(all_success_times),
            key3='百分比(评分)',
            value3=str(round((all_success_times * 100/all_times), 3) if all_times else 0),
            status=BACKTESTSTATUS.SUCCESS
        )

        # 更新数据库
        db.Backtest.set(**res).where(id=backtest_id)

        return res
    except Exception as e:
        err = traceback.format_exc()
        db.Backtest.set(status=BACKTESTSTATUS.ERROR, mark=str(err)).where(id=backtest_id)



class BacktestController(object):
    """ stock """

    def __init__(self, id=None):
        self.id = id

    @property
    def backtest(self):
        return db.Backtest.get(id=self.id)

    @property
    def record(self):
        return db.BacktestRecord.find(backtest_id=self.id)

    @classmethod
    def income_times(cls, args):
        """ 策略的收益次数统计 """
        # 查找股票
        p=Process(
            target=income_times_run,
            name="income_times_run",
            args=(
                args,
                args["strategy_name"],
                args["from_date_int"],
                args["low_price"],
                args["high_price"],
                args["days"],
                args["percent"]
            )
        )
        p.deamon = False
        p.start()
        return p.pid

    @classmethod
    def all_backtest(cls):
        return db.Backtest.select().order_by("id desc").all()

    def kill_run(self):
        """结束一个进程"""
        pid = self.backtest.get("pid")
        if pid:
            kill_pid(int(pid))
            db.Backtest.set(status=BACKTESTSTATUS.ERROR).where(id=self.id)

    def delete(self):
        self.kill_run()
        db.BacktestRecord.rm(backtest_id=self.id)
        db.Backtest.rm(id=self.id)


